# 主程序集成模块
from DrissionPage import WebPage, ChromiumOptions
import json
import os
import time
import random
from collect_data import DataCollector
from collect_comments import CommentCollector
from collect_notes import NotesCollector

# 提供的小红书cookies
cookies = ['abRequestId=ea1b36e3-290b-5589-800d-ef6e4b7fed88; xsecappid=xhs-pc-web; a1=196ab51cec7t1ks552aduq8k2e1y9ed7ewbub2fi130000329117; webId=2ffc3ee8a1dd24d6d32f8d7a248d2307; acw_tc=0a0bb4b917466309878964371e60dd118106e99a359d579c40416d2dd21f0c; gid=yjK0D2yfSjd8yjK0D2ySdEJiSW9yTM22J0f7D3yAYTJdy7q8vjdfWD888qJjyyW8jWjSD8jK; webBuild=4.62.3; websectiga=2845367ec3848418062e761c09db7caf0e8b79d132ccdd1a4f8e64a11d0cac0d; sec_poison_id=e581ef89-e70f-44aa-988e-b4400f5a96ad; loadts=1746631540844; unread={%22ub%22:%2267f8726c000000001c0158ec%22%2C%22ue%22:%2267f993cb000000001d01c3b4%22%2C%22uc%22:29}; web_session=040069b45d0634f8cf12fb332e3a4bce08a203']

class XHSCrawler:
    def __init__(self, port=None):
        # 创建浏览器选项实例
        co = ChromiumOptions()
        
        # 设置随机User-Agent
        co.set_user_agent(self.get_random_user_agent())
        
        # 设置浏览器指纹
        co.set_argument('--disable-blink-features=AutomationControlled')  # 禁用自动化控制特征
        co.set_argument('--disable-infobars')  # 禁用信息栏
        co.set_argument('--disable-notifications')  # 禁用通知
        co.set_argument('--disable-popup-blocking')  # 禁用弹窗拦截
        co.set_argument('--disable-save-password-bubble')  # 禁用保存密码提示
        co.set_argument('--disable-translate')  # 禁用翻译
        co.set_argument('--disable-web-security')  # 禁用网页安全策略
        co.set_argument('--no-sandbox')  # 禁用沙箱
        co.set_argument('--disable-gpu')  # 禁用GPU加速
        co.set_argument('--disable-dev-shm-usage')  # 禁用/dev/shm使用
        co.set_argument('--disable-extensions')  # 禁用扩展
        co.set_argument('--disable-plugins')  # 禁用插件
        co.set_argument('--disable-software-rasterizer')  # 禁用软件光栅化
        co.set_argument('--disable-features=IsolateOrigins,site-per-process')  # 禁用站点隔离
        co.set_argument('--disable-site-isolation-trials')  # 禁用站点隔离试验
        co.set_argument('--disable-web-security')  # 禁用网页安全策略
        co.set_argument('--disable-features=BlockInsecurePrivateNetworkRequests')  # 禁用不安全私有网络请求拦截
        co.set_argument('--disable-features=IsolateOrigins')  # 禁用源隔离
        co.set_argument('--disable-features=site-per-process')  # 禁用站点进程隔离
        
        # 设置初始窗口大小
        co.set_argument('--window-size', '1200,1000')
        
        # 如果指定了端口，则设置端口
        if port:
            co.set_port(port)
            
        # 创建浏览器实例
        self.page = WebPage(chromium_options=co)
        
        # 注入反检测JS
        self.page.run_js("""
            Object.defineProperty(navigator, 'webdriver', {
                get: () => undefined
            });
            Object.defineProperty(navigator, 'plugins', {
                get: () => [1, 2, 3, 4, 5]
            });
            Object.defineProperty(navigator, 'languages', {
                get: () => ['zh-CN', 'zh', 'en']
            });
        """)
        
        # 创建采集器实例
        self.comment_collector = CommentCollector()
        self.notes_collector = NotesCollector()
        self.data_collector = DataCollector()  # 添加DataCollector实例
        
        # 添加关键词属性
        self.current_keyword = ""
        
        # 添加新的属性来存储采集数量限制
        self.max_notes = 20  # 默认值
        self.max_comments = 20  # 默认值
        self.cookie_path = None  # 添加cookie路径属性
    
    def get_random_user_agent(self):
        """生成随机User-Agent"""
        user_agents = [
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/121.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/120.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/119.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/122.0.0.0 Safari/537.36",
            "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/605.1.15 (KHTML, like Gecko) Version/17.3 Safari/605.1.15"
        ]
        return random.choice(user_agents)
    
    # 添加新的方法来设置采集数量
    def set_limits(self, max_notes=20, max_comments=20):
        """设置采集数量限制"""
        self.max_notes = max_notes
        self.max_comments = max_comments
        print(f"设置采集限制：每个关键词{max_notes}篇笔记，每篇笔记{max_comments}条评论")

    def set_cookie_path(self, cookie_path):
        """设置cookie文件路径"""
        self.cookie_path = cookie_path

    def load_cookies(self):
        """加载指定的cookies文件"""
        try:
            if not self.cookie_path:
                # 如果没有指定cookie路径，使用默认的cookies目录
                cookies_dir = 'cookies'
                if not os.path.exists(cookies_dir):
                    raise FileNotFoundError("cookies目录不存在")
                
                cookie_files = [f for f in os.listdir(cookies_dir) if f.endswith('.json')]
                if not cookie_files:
                    raise FileNotFoundError("没有找到cookies文件")
                
                latest_file = max(cookie_files, 
                                key=lambda x: os.path.getctime(os.path.join(cookies_dir, x)))
                cookie_path = os.path.join(cookies_dir, latest_file)
            else:
                cookie_path = self.cookie_path
            
            # 读取cookies
            with open(cookie_path, 'r', encoding='utf-8') as f:
                return json.load(f)
                
        except Exception as e:
            print(f"加载cookies失败: {e}")
            return None

    def login(self):
        """使用cookies登录小红书"""
        try:
            # 先访问小红书
            self.page.get('https://www.xiaohongshu.com')
            # 随机等待时间，模拟真实用户行为
            time.sleep(random.uniform(2, 4))
            
            # 设置cookies
            self.page.set.cookies(cookies)
            print("已添加小红书cookie")
            
            # 刷新页面使cookies生效
            self.page.refresh()
            # 随机等待时间，模拟真实用户行为
            time.sleep(random.uniform(3, 5))
            
            # 检查登录状态
            print("成功导航到小红书并应用Cookie！")
            return True
            
        except Exception as e:
            print(f"登录失败: {e}")
            return False

    def search_and_filter(self, keyword=''):
        """搜索关键词并筛选"""
        try:
            # 保存当前关键词
            self.current_keyword = keyword
            
            # 设置当前关键词到采集器
            self.comment_collector.set_keyword(keyword)
            self.notes_collector.set_keyword(keyword)
            self.data_collector.set_keyword(keyword)  # 也设置到DataCollector
            
            # 定位搜索框并输入内容
            search_input = self.page.ele('#search-input')
            search_input.input(f'{keyword}\n')
            time.sleep(2)  # 增加等待时间
            
            # 定位filter区域，执行鼠标悬浮
            filter_area = self.page.ele('.filter')
            filter_area.hover()
            time.sleep(2)  # 增加等待时间
            
            # 使用多种选择器尝试定位"最热"选项
            # 方法1: 使用固定XPath
            hot_tab = self.page.ele('xpath:/html/body/div[2]/div[1]/div[2]/div[2]/div/div/div[1]/div[2]/div/div/div[1]/div[1]/div/div[3]')
            
            # 方法2: 如果方法1失败，尝试通过文本内容查找
            if not hot_tab:
                hot_tab = self.page.ele('xpath://div[contains(text(), "最多点赞")]')
                
            # 方法3: 尝试找到排序选项的父元素，然后选择第3个子元素(通常是"最热")
            if not hot_tab:
                filter_tabs = self.page.eles('.filter div[role="button"]')
                if len(filter_tabs) >= 3:
                    hot_tab = filter_tabs[2]  # 假设第3个是"最热"
            
            # 方法4: 尝试通过CSS选择器查找
            if not hot_tab:
                hot_tab = self.page.ele('.filter div:nth-child(3)')
                
            # 如果找到了最热选项，点击它
            if hot_tab:
                hot_tab.click()
                print("已点击最热选项")
                time.sleep(2)
                return True
            else:
                print("未找到最热选项，将继续使用默认排序")
                # 即使找不到最热选项，也返回True以继续执行
                return True
                
        except Exception as e:
            print(f"搜索筛选失败: {e}")
            # 由于可能是界面结构变化导致的，我们返回True以继续执行后续流程
            return True

    def collect_note_and_comments(self, note_element):
        """采集单篇笔记及其评论"""
        try:
            # 增加点击前的等待时间
            time.sleep(2)
            note_element.click()
            time.sleep(2)  # 增加点击后的等待时间
            
            # 设置当前关键词到数据采集器
            self.data_collector.set_keyword(self.current_keyword)
            
            # 使用类的数据采集器实例，而不是创建新的
            if self.data_collector.collect_article_data(self.page):
                self.comment_collector.processed_comments.clear()
                
                comment_count = 0
                no_new_comments_count = 0
                max_attempts = 2  # 最大尝试次数
                attempts = 0
                
                while comment_count < self.max_comments:
                    time.sleep(1)  # 每次循环前等待
                    comments_container = self.page.ele('.comments-el')
                    if not comments_container:
                        print("未找到评论容器")
                        break  # 直接跳出循环，继续下一篇笔记
                    
                    # 获取滚动前的评论元素数量
                    current_comments = self.page.eles('.list-container')
                    previous_count = len(current_comments) if current_comments else 0
                    
                    # 滚动加载更多评论
                    self.page.actions.scroll(delta_y=1800)
                    time.sleep(1.5)  # 等待新内容加载
                    
                    # 获取滚动后的评论元素数量
                    new_comments = self.page.eles('.comment-item')
                    new_count = len(new_comments) if new_comments else 0
                    
                    # 检查是否有新评论元素加载
                    if new_count <= previous_count:
                        no_new_comments_count += 1
                        print(f"连续{no_new_comments_count}次未发现新评论元素")
                    else:
                        no_new_comments_count = 0
                    
                    # 采集当前可见的评论
                    if self.comment_collector.collect_comments(self.page):
                        comment_count = len(self.comment_collector.processed_comments)
                        print(f"当前已采集{comment_count}条评论")
                    
                    # 增加尝试次数
                    attempts += 1
                    
                    # 如果连续没有新评论或者尝试次数达到上限，且评论数量少于目标值，则退出
                    if (no_new_comments_count >= 2 or attempts >= max_attempts) and comment_count < self.max_comments:
                        actual_total = new_count if new_count > 0 else comment_count
                        print(f"笔记评论可能不足{self.max_comments}条，当前可见评论数为{actual_total}条，继续下一篇笔记")
                        break
                    
                    # 如果还没有任何评论且已尝试多次，则退出
                    if comment_count == 0 and attempts >= 2:
                        print("该笔记可能没有评论，继续下一篇笔记")
                        break
                    
                    # 如果已达到目标评论数，退出循环
                    if comment_count >= self.max_comments:
                        break
                
                print(f"笔记评论采集完成，共采集{comment_count}条评论")
                return True  # 即使没有达到目标评论数也返回成功
            
            return False
            
        except Exception as e:
            print(f"采集笔记和评论时出错: {e}")
            return False
        finally:
            time.sleep(1)  # 返回前等待
            self.page.back()
            time.sleep(2)  # 增加返回后的等待时间

    def collect_all_notes(self):
        """遍历采集笔记和评论"""
        try:
            collected_count = 0
            no_new_notes_count = 0
            processed_indexes = set()  # 记录已处理的笔记索引
            
            while collected_count < self.max_notes:
                time.sleep(1.5)  # 每次循环前等待
                
                # 获取当前可见的笔记
                note_items = self.page.eles('.note-item')
                current_notes_count = len(note_items) if note_items else 0
                
                if not note_items:
                    print("未找到笔记项")
                    if no_new_notes_count >= 2:
                        break
                    no_new_notes_count += 1
                    time.sleep(2)  # 未找到笔记时等待更长时间
                    continue
                
                # 遍历当前可见的所有笔记
                for i in range(current_notes_count):
                    if collected_count >= self.max_notes:
                        print(f"已达到{self.max_notes}篇笔记上限")
                        return True
                    
                    # 跳过已处理的笔记
                    if i in processed_indexes:
                        continue
                    
                    try:
                        time.sleep(1)  # 处理每个笔记前等待
                        
                        # 重新获取最新的笔记列表和当前要处理的笔记
                        current_notes = self.page.eles('.note-item')
                        if i >= len(current_notes):
                            break
                        
                        current_note = current_notes[i]
                        if self.collect_note_and_comments(current_note):
                            collected_count += 1
                            processed_indexes.add(i)  # 记录已处理的索引
                            print(f"已采集第{collected_count}篇笔记及其评论")
                            no_new_notes_count = 0
                            time.sleep(2)  # 成功采集后等待
                        
                    except Exception as e:
                        print(f"处理单个笔记时出错: {e}")
                        time.sleep(3)  # 出错后等待更长时间
                        continue
                
                # 如果当前页面的笔记都处理完了，继续滚动加载更多
                if len(processed_indexes) >= current_notes_count:
                    # 滚动加载更多笔记
                    self.page.actions.scroll(delta_y=2300)
                    time.sleep(3)  # 增加滚动后的等待时间
                    
                    # 获取滚动后的笔记数量
                    new_notes = self.page.eles('.note-item')
                    if len(new_notes) <= current_notes_count:
                        no_new_notes_count += 1
                        if no_new_notes_count >= 2:
                            print("已无法加载更多笔记")
                            break
                    else:
                        no_new_notes_count = 0
            
            return True
            
        except Exception as e:
            print(f"遍历采集笔记出错: {e}")
            return False

    def close(self):
        """关闭浏览器"""
        if hasattr(self, 'page'):
            self.page.quit()

def main():
    crawler = XHSCrawler()
    try:
        if crawler.login() and \
           crawler.search_and_filter() and \
           crawler.collect_all_notes():
            print("所有数据采集完成")
        else:
            print("数据采集过程中出现错误")
    finally:
        input("按回车键退出...")
        crawler.close()

if __name__ == '__main__':
    main()
